{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b50c4af8-fec3-4396-860a-1322089d76cb",
   "metadata": {},
   "source": [
    "# ReAct Agent with Query Engine Tools"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db402a8b-90d6-4e1d-8df6-347c54624f26",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Build Query Engine Tools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "02160804-64a2-4ef3-8a0d-8c16b06fd205",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from llama_index import (\n",
    "    SimpleDirectoryReader,\n",
    "    VectorStoreIndex,\n",
    "    StorageContext,\n",
    "    load_index_from_storage,\n",
    ")\n",
    "\n",
    "from llama_index.tools import QueryEngineTool, ToolMetadata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "91618236-54d3-4783-86b7-7b7554efeed1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "try:\n",
    "    storage_context = StorageContext.from_defaults(persist_dir=\"./storage/lyft\")\n",
    "    lyft_index = load_index_from_storage(storage_context)\n",
    "\n",
    "    storage_context = StorageContext.from_defaults(persist_dir=\"./storage/uber\")\n",
    "    uber_index = load_index_from_storage(storage_context)\n",
    "\n",
    "    index_loaded = True\n",
    "except:\n",
    "    index_loaded = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d3d0bb8c-16c8-4946-a9d8-59528cf3952a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "if not index_loaded:\n",
    "    # load data\n",
    "    lyft_docs = SimpleDirectoryReader(\n",
    "        input_files=[\"../data/10k/lyft_2021.pdf\"]\n",
    "    ).load_data()\n",
    "    uber_docs = SimpleDirectoryReader(\n",
    "        input_files=[\"../data/10k/uber_2021.pdf\"]\n",
    "    ).load_data()\n",
    "\n",
    "    # build index\n",
    "    lyft_index = VectorStoreIndex.from_documents(lyft_docs)\n",
    "    uber_index = VectorStoreIndex.from_documents(uber_docs)\n",
    "\n",
    "    # persist index\n",
    "    lyft_index.storage_context.persist(persist_dir=\"./storage/lyft\")\n",
    "    uber_index.storage_context.persist(persist_dir=\"./storage/uber\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "31892898-a2dc-43c8-812a-3442feb2108d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "lyft_engine = lyft_index.as_query_engine(similarity_top_k=3)\n",
    "uber_engine = uber_index.as_query_engine(similarity_top_k=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "f9f3158a-7647-4442-8de1-4db80723b4d2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "query_engine_tools = [\n",
    "    QueryEngineTool(\n",
    "        query_engine=lyft_engine,\n",
    "        metadata=ToolMetadata(\n",
    "            name=\"lyft_10k\",\n",
    "            description=\"Provides information about Lyft financials for year 2021. \"\n",
    "            \"Use a detailed plain text question as input to the tool.\",\n",
    "        ),\n",
    "    ),\n",
    "    QueryEngineTool(\n",
    "        query_engine=uber_engine,\n",
    "        metadata=ToolMetadata(\n",
    "            name=\"uber_10k\",\n",
    "            description=\"Provides information about Uber financials for year 2021. \"\n",
    "            \"Use a detailed plain text question as input to the tool.\",\n",
    "        ),\n",
    "    ),\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "275c01b1-8dce-4216-9203-1e961b7fc313",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Setup ReAct Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "32f71a46-bdf6-4365-b1f1-e23a0d913a3d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from llama_index.agent import ReActAgent\n",
    "from llama_index.llms import OpenAI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ded93297-fee8-4329-bf37-cf77e87621ae",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "llm = OpenAI(model=\"gpt-3.5-turbo-0613\")\n",
    "agent = ReActAgent.from_tools(query_engine_tools, llm=llm, verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "70a82471-9226-42ad-bd8a-aebde3530d95",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[38;5;200m\u001b[1;3mThought: I need to use a tool to help me answer the question.\n",
      "Action: lyft_10k\n",
      "Action Input: {'input': \"What was Lyft's revenue growth in 2021?\"}\n",
      "\u001b[0m\u001b[36;1m\u001b[1;3mObservation: Lyft's revenue growth in 2021 was 36%.\n",
      "\u001b[0m\u001b[38;5;200m\u001b[1;3mResponse: Lyft's revenue growth in 2021 was 36%.\n",
      "\u001b[0mLyft's revenue growth in 2021 was 36%.\n"
     ]
    }
   ],
   "source": [
    "response = agent.chat(\"What was Lyft's revenue growth in 2021?\")\n",
    "print(str(response))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "0873463d-4790-4c17-bfe9-2ece610fe4b3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[38;5;200m\u001b[1;3mThought: I need to use a tool to help me compare the revenue growth of Uber and Lyft in 2021.\n",
      "Action: lyft_10k\n",
      "Action Input: {'input': \"What was Lyft's revenue growth in 2021?\"}\n",
      "\u001b[0m\u001b[36;1m\u001b[1;3mObservation: Lyft's revenue growth in 2021 was 36%.\n",
      "\u001b[0m\u001b[38;5;200m\u001b[1;3mThought: I need to use a tool to help me compare the revenue growth of Uber and Lyft in 2021.\n",
      "Action: uber_10k\n",
      "Action Input: {'input': \"What was Uber's revenue growth in 2021?\"}\n",
      "\u001b[0m\u001b[36;1m\u001b[1;3mObservation: Uber's revenue growth in 2021 was 57%.\n",
      "\u001b[0m\u001b[38;5;200m\u001b[1;3mResponse: In 2021, Lyft's revenue growth was 36% while Uber's revenue growth was 57%. This indicates that Uber experienced a higher revenue growth compared to Lyft in 2021.\n",
      "\u001b[0mIn 2021, Lyft's revenue growth was 36% while Uber's revenue growth was 57%. This indicates that Uber experienced a higher revenue growth compared to Lyft in 2021.\n"
     ]
    }
   ],
   "source": [
    "response = agent.chat(\n",
    "    \"Compare and contrast the revenue growth of Uber and Lyft in 2021, then give an analysis\"\n",
    ")\n",
    "print(str(response))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb63492a-836c-42da-94a4-0b22cccbc3e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Try another query with async execution\n",
    "\n",
    "import nest_asyncio\n",
    "\n",
    "nest_asyncio.apply()\n",
    "\n",
    "response = await agent.achat(\n",
    "    \"Compare and contrast the risks of Uber and Lyft in 2021, then give an analysis\"\n",
    ")\n",
    "print(str(response))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
